package com.shanks.framework.mq.rocketmq.listener;

import com.shanks.framework.mq.common.constant.MQConstant;
import com.shanks.framework.mq.common.util.MDCUtil;
import com.shanks.framework.mq.domain.MQMessage;
import com.shanks.framework.mq.processor.MQProcessor;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyContext;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.common.message.MessageExt;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.UnsupportedEncodingException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 封装RocketMQ并发消费
 */
public class ConcurrentMessageListener implements MessageListenerConcurrently {
    private static final Logger logger = LoggerFactory.getLogger(ConcurrentMessageListener.class);

    private Map<String, MQProcessor> mqProcessorMap = new HashMap<>();

    public void addMQProcessor(String topic, MQProcessor messageProcessor) {
        mqProcessorMap.put(topic, messageProcessor);
    }

    public MQProcessor getByTopic(String topic) {
        return mqProcessorMap.get(topic);
    }

    @Override
    public ConsumeConcurrentlyStatus consumeMessage(List<MessageExt> messageExtList, ConsumeConcurrentlyContext consumeConcurrentlyContext) {
        try {
            MDCUtil.putTraceId();
            for (MessageExt msg : messageExtList) {
                try {
                    MQMessage message = convertMessageWithProps(msg);
                    boolean result = getByTopic(message.getTopic()).handleMessage(message);
                    if (!result)
                        return ConsumeConcurrentlyStatus.RECONSUME_LATER;
                    logger.info("处理MQ消息成功, topic:{}, tags:{}, msgId:{}", msg.getTopic(), msg
                            .getTags(), msg.getMsgId());
                } catch (Throwable e) {
                    logger.error("处理MQ消息失败, topic:{}, tags:{}, msgId:{}, body:{}, error:{}", msg.getTopic(), msg
                            .getTags(), msg.getMsgId(), getBodyAsString(msg.getBody()), e.getMessage(), e);
                    return ConsumeConcurrentlyStatus.RECONSUME_LATER;
                }
            }
            return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
        } catch (Exception e) {
            throw e;
        } finally {
            MDCUtil.removeTraceId();
        }
    }

    /**
     * 转换成自定义的消费消息体
     *
     * @param msg
     * @return
     */
    private MQMessage convertMessageWithProps(MessageExt msg) {
        Map<String, Object> props = new HashMap<>();
        props.put(MQConstant.MSG.MSG_SOURCE, msg.getUserProperty(MQConstant.MSG.MSG_SOURCE));
        props.put(MQConstant.MSG.PROPERTY_SIGNATURE, MQConstant.MSG.PROPERTY_SIGNATURE);
        MQMessage mqMessage = new MQMessage(msg.getMsgId(), msg.getTopic(), msg.getTags(), msg.getKeys(), msg.getBody(), msg, props);
        return mqMessage;
    }

    private String getBodyAsString(byte[] body) {
        try {
            return new String(body, "UTF-8");
        } catch (UnsupportedEncodingException e) {
            return null;
        }
    }
}
